from typing import List, Optional, Tuple

import torch
import torch.nn as nn
from mmcv.ops import batched_nms
from mmengine.config import ConfigDict
from mmengine.model import bias_init_with_prob, normal_init
from mmengine.structures import InstanceData
from torch import Tensor

from mmdet.registry import MODELS
from mmdet.utils import (ConfigType, InstanceList, OptConfigType,
                         OptInstanceList, OptMultiConfig)
from ..utils import (gaussian_radius, gen_gaussian_target, get_local_maximum,
                     get_topk_from_heatmap, multi_apply,
                     transpose_and_gather_feat)
from .base_dense_head import BaseDenseHead
from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule
@MODELS.register_module()
class Slot_head(BaseDenseHead):
    def __init__(self,
                 in_channels: list,
                 feat_channels: list,
                 num_classes:int,
                 stack_nums:list,
                 norm_cfg: ConfigType = dict(type='BN', momentum=0.03, eps=0.001),
                 act_cfg: ConfigType = dict(type='Swish'),
                 use_depthwise = False,
                 loss_center_heatmap: ConfigType = dict(
                     type='GaussianFocalLoss', loss_weight=5.0),
                 loss_cls: ConfigType = dict(type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
                 loss_obj: ConfigType = dict(type='CrossEntropyLoss', use_sigmoid=True, loss_weight=5.0),
                 loss_ori: ConfigType = dict(type='L1Loss', loss_weight=1.0),
                 loss_offset: ConfigType = dict(type='L1Loss', loss_weight=1.0),
                 train_cfg: OptConfigType = None,
                 test_cfg: OptConfigType = None,
                 init_cfg: OptMultiConfig = None) -> None:
        super().__init__(init_cfg=init_cfg)
        self.in_channels = in_channels
        self.feat_channels = feat_channels
        self.num_classes = num_classes
        self.test_cfg = test_cfg
        self.loss_center_heatmap = MODELS.build(loss_center_heatmap)
        self.loss_cls =  MODELS.build(loss_cls)
        self.loss_offset = MODELS.build(loss_offset)

        self.loss_obj = MODELS.build(loss_obj)
        self.loss_ori = MODELS.build(loss_ori)

        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
        self.stack_nums = stack_nums
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg
        self.use_depthwise = use_depthwise
        self.init_Head()
    def init_Head(self):
        self.joint_stem = nn.Sequential()
        self.slot_stem = nn.Sequential()

        conv = DepthwiseSeparableConvModule \
            if self.use_depthwise else ConvModule
        for i in range(self.stack_nums[0]):
            chn = self.in_channels[0] if i == 0 else self.feat_channels[0]
            self.joint_stem.append(
                conv(
                    chn,
                    self.feat_channels[0],
                    3,
                    stride=1,
                    padding=1,
                    norm_cfg=self.norm_cfg,
                    act_cfg=self.act_cfg,
                    ))
        
        for i in range(self.stack_nums[1]):
            chn = self.in_channels[0] if i == 0 else self.feat_channels[1]
            self.slot_stem.append(
                conv(
                    chn,
                    self.feat_channels[1],
                    3,
                    stride=1,
                    padding=1,
                    norm_cfg=self.norm_cfg,
                    act_cfg=self.act_cfg,
                    )
            )
        self.conv_joint_obj = nn.Conv2d(self.feat_channels[0],1,3,padding=1)
        self.conv_joint_offset = nn.Conv2d(self.feat_channels[0],2,3,padding=1)
        # self.conv_joint_ori = nn.Conv2d(self.feat_channels[0],2,3,padding=1)

        self.conv_slot_cls = nn.Conv2d(self.feat_channels[1],self.num_classes,3,padding=1)
        self.conv_slot_offset = nn.Conv2d(self.feat_channels[1],4,3,padding=1)
        self.conv_slot_ori = nn.Conv2d(self.feat_channels[1],2,3,padding=1)
        self.conv_slot_ct = nn.Conv2d(self.feat_channels[1],2,3,padding=1)

    def forward(self,input_tensors):
        assert len(input_tensors) == 2
        joint_feature = self.joint_stem(input_tensors[0])
        slot_feature = self.slot_stem(input_tensors[1])

        joint_obj = self.conv_joint_obj(joint_feature)
        joint_offset = self.conv_joint_offset(joint_feature)
        # joint_ori = self.conv_joint_ori(joint_feature)

        slot_cls = self.conv_slot_cls(slot_feature)
        slot_offset = self.conv_slot_offset(slot_feature)
        slot_ori = self.conv_slot_ori(slot_feature)
        slot_ct = self.conv_slot_ct(slot_feature)
        return (slot_cls,slot_offset,slot_ori,slot_ct,joint_obj,joint_offset)
    def loss_by_feat():
        return 0
    def loss(self, features,gt_datas):
        slot_cls,slot_offset,slot_ori,slot_ct,joint_obj,joint_offset = self.forward(features)
        gt_joints  = [torch.tensor(gt_data.mark_points,device=slot_cls.device,dtype=torch.float32).reshape(-1,2)  for gt_data in gt_datas]
        gt_slots = [ torch.tensor(gt_data.slots,device=slot_cls.device,dtype=torch.float32).reshape(-1,8) for gt_data in gt_datas]
        img_shape = gt_datas[0].img_shape



        target_result = self.get_targets(gt_joints,gt_slots,[features[0].shape,features[1].shape],img_shape)
        center_heatmap_target = target_result['center_heatmap_target']
        slot_offset_target = target_result['slot_offset_target']
        slot_ori_target = target_result['slot_ori_target']
        slot_ct_target = target_result['slot_ct_target']
        slot_weight_target = target_result['slot_weight_target']

        obj_target = target_result['obj_target']
        joint_offset_target = target_result['joint_offset_target']

     
        avg_factor1 = target_result['avg_factor1']
        avg_factor2 = target_result['avg_factor2']

        # img = gt_datas[1].img
        # import cv2
        # import numpy as np
        # obj_target_vis = cv2.resize((obj_target[1]*255).cpu().numpy().transpose(1,2,0).astype(np.uint8),(512,512))
        # center_heatmap_target_vis = cv2.resize((center_heatmap_target[1]*255).cpu().numpy().transpose(1,2,0).astype(np.uint8),(512,512))
        # img = img + obj_target_vis.reshape(512,512,1) + center_heatmap_target_vis
        # cv2.imshow('img0',img)
        # # cv2.imshow('heatmap',center_heatmap_target_vis) 
        # cv2.waitKey()



        loss_center_heatmap = self.loss_center_heatmap(slot_cls.sigmoid(),center_heatmap_target,avg_factor=slot_cls.shape[2]*slot_cls.shape[3])
        loss_slot_offset = self.loss_offset(slot_offset,slot_offset_target,slot_weight_target,avg_factor = avg_factor1)
        loss_slot_ori = self.loss_offset(slot_ori,slot_ori_target,slot_weight_target,avg_factor = avg_factor1)
        loss_slot_ct = self.loss_offset(slot_ct,slot_ct_target,slot_weight_target,avg_factor = avg_factor1)

        loss_joint_obj = self.loss_obj(joint_obj,obj_target,avg_factor=joint_obj.shape[2]*joint_obj.shape[3])
        loss_joint_offset = self.loss_offset(joint_offset,joint_offset_target,obj_target,avg_factor = avg_factor2)

        return dict(
            loss_center_heatmap = loss_center_heatmap,
            loss_slot_ct = loss_slot_ct,
            loss_slot_offset = loss_slot_offset,
            loss_slot_ori = loss_slot_ori,
            loss_joint_obj = loss_joint_obj,
            loss_joint_offset = loss_joint_offset
        ) 

    def predict(self,x, batch_data_samples, rescale=False):
        img_shape = batch_data_samples[0].img_shape
        img_h,img_w = img_shape
        

        slot_cls,slot_offset,slot_ori,slot_ct,joint_obj,joint_offset = self.forward(x)
        slot_cls = slot_cls.sigmoid()
        joint_obj = joint_obj.sigmoid()
        # 处理slot
        slot_feat_shape = slot_cls.shape[2:]
        feat_h,feat_w = slot_feat_shape

        h_ratio = img_h /feat_h
        w_ratio = img_w /feat_w

        

        slot_cls = get_local_maximum(slot_cls,kernel=3)
        *batch_dets, topk_ys, topk_xs = get_topk_from_heatmap(
            slot_cls, k=50)
        batch_scores, batch_index, batch_topk_labels = batch_dets

        slot_offset = transpose_and_gather_feat(slot_offset,batch_index)
        slot_ori = transpose_and_gather_feat(slot_ori,batch_index)
        slot_ct = transpose_and_gather_feat(slot_ct,batch_index)

        center_x = topk_xs + slot_ct[...,0]
        center_y = topk_ys + slot_ct[...,1]
# TODO: change offset predict
        # slot_offset1_x = center_x + slot_offset[...,0]
        # slot_offset1_y = center_y + slot_offset[...,1]

        # slot_offset2_x = center_x + slot_offset[...,2]
        # slot_offset2_y = center_y + slot_offset[...,3]
        slot_offset1_x =slot_offset[...,0]
        slot_offset1_y =slot_offset[...,1]

        slot_offset2_x =slot_offset[...,2]
        slot_offset2_y =slot_offset[...,3]

        ratio = torch.tensor([w_ratio,h_ratio],dtype=slot_cls.dtype,device=slot_cls.device)

        slot_center = torch.concat([center_x.unsqueeze(2),center_y.unsqueeze(2)],dim=-1)*ratio
        slot_joint_l = torch.concat([slot_offset1_x.unsqueeze(2),slot_offset1_y.unsqueeze(2)],dim=-1)*ratio
        slot_joint_r = torch.concat([slot_offset2_x.unsqueeze(2),slot_offset2_y.unsqueeze(2)],dim=-1)*ratio
        predict_labels = []
        predict_scores = []
        predict_centers = []
        predict_joint_l = []
        predict_joint_r = []
        predict_ori = []
        bs = batch_topk_labels.shape[0]
        for batch_id in range(bs):
            index0 = batch_scores[batch_id,:]>self.test_cfg['slot_entrance_thr']
            
            score = batch_scores[batch_id,index0].cpu().numpy()
            label = batch_topk_labels[batch_id,index0].cpu().numpy()
            center = slot_center[batch_id,index0,:].cpu().numpy()
            joint_l = slot_joint_l[batch_id,index0,:].cpu().numpy()
            joint_r = slot_joint_r[batch_id,index0,:].cpu().numpy()
            ori = slot_ori[batch_id,index0,:].cpu().numpy()

            predict_labels.append(label)
            predict_scores.append(score)
            predict_centers.append(center)
            predict_joint_l.append(joint_l)
            predict_joint_r.append(joint_r)
            predict_ori.append(ori)

        # 处理joint
        joint_feat_shape = joint_obj.shape[2:]
        feat_h,feat_w = joint_feat_shape

        h_ratio = img_h /feat_h
        w_ratio = img_w /feat_w
        joint_obj = get_local_maximum(joint_obj,kernel=5)
        *batch_dets, topk_ys, topk_xs = get_topk_from_heatmap(
            joint_obj, k=50)
        batch_scores, batch_index, batch_topk_labels = batch_dets
        joint_offset = transpose_and_gather_feat(joint_offset,batch_index)

        joint_x = topk_xs + joint_offset[...,0]
        joint_y = topk_ys + joint_offset[...,1]

        ratio = torch.tensor([w_ratio,h_ratio],dtype=slot_cls.dtype,device=slot_cls.device)

        joint_center = torch.concat([joint_x.unsqueeze(2),joint_y.unsqueeze(2)],dim=-1)*ratio
        result = InstanceData()

        

        predict_joint_score = []
        predict_joint_center = []
        for batch_id in range(bs):
            index1 = batch_scores[batch_id]>self.test_cfg["joint_thr"]
            score = batch_scores[batch_id,index1].cpu().numpy()
            center = joint_center[batch_id,index1,:].cpu().numpy()
            predict_joint_score.append(score)
            predict_joint_center.append(center)




        result.labels = predict_labels
        result.scores = predict_scores
        result.center = predict_centers
        result.joint_l = predict_joint_l
        result.joint_r = predict_joint_r
        result.ori = predict_ori
        result.joint_obj = predict_joint_score
        result.joint_center = predict_joint_center

        return result
        




    def get_targets(self,gt_joints,gt_slots,feat_shape,img_shape):
        img_h,img_w = img_shape
        # 处理slot对应的真值
        bs,_,feat_h,feat_w = feat_shape[1]
        width_ratio = float(feat_w/img_w)
        height_ratio = float(feat_h/img_h)

        center_heatmap_target = gt_slots[-1].new_zeros(bs,self.num_classes,feat_h,feat_w)
        slot_offset_target = gt_slots[-1].new_zeros(bs,4,feat_h,feat_w)
        slot_ori_target = gt_slots[-1].new_zeros(bs,2,feat_h,feat_w)
        slot_ct_target = gt_slots[-1].new_zeros(bs,2,feat_h,feat_w)

        slot_weight_target = gt_slots[-1].new_zeros(bs,1,feat_h,feat_w)
        for batch_id in range(bs):
            gt_slot = gt_slots[batch_id]
            center_x = gt_slot[:,0].reshape(-1,1) * width_ratio
            center_y = gt_slot[:,1].reshape(-1,1) * height_ratio
            gt_centers = torch.cat((center_x, center_y), dim=-1)
            for j, ct in enumerate(gt_centers):
                ctx_int,cty_int = ct.int()
                ctx,cty = ct

                r1 = torch.sqrt(gt_slot[j,2]**2 + gt_slot[j,3]**2) 
                r2 = torch.sqrt(gt_slot[j,4]**2 + gt_slot[j,5]**2) 
                radius = torch.round(min(r1*width_ratio,r2*width_ratio)).long()
                ind = gt_slot[j][7].long()

                gen_gaussian_target(center_heatmap_target[batch_id, ind],
                                    [ctx_int, cty_int], radius)
                slot_offset_target[batch_id,:,cty_int,ctx_int] = gt_slot[j,2:6]*width_ratio
                slot_ori_target[batch_id,0,cty_int,ctx_int] = torch.sin(gt_slot[j,6])
                slot_ori_target[batch_id,1,cty_int,ctx_int] = torch.cos(gt_slot[j,6])

                slot_ct_target[batch_id,0,cty_int,ctx_int] = ctx - ctx_int
                slot_ct_target[batch_id,1,cty_int,ctx_int] = cty - cty_int

                slot_weight_target[batch_id,:,cty_int,ctx_int] = 1
        # 处理joint对应的gt
        bs,_,feat_h,feat_w = feat_shape[0]
        width_ratio = float(feat_w/img_w)
        height_ratio = float(feat_h/img_h)

        obj_target =  gt_joints[-1].new_zeros(bs,1,feat_h,feat_w)
        joint_offset_target =  gt_joints[-1].new_zeros(bs,2,feat_h,feat_w)
        for batch_id in range(bs):
            gt_joint = gt_joints[batch_id]
            center_x_joint = gt_joint[:,0].reshape(-1,1) * width_ratio
            center_y_joint = gt_joint[:,1].reshape(-1,1) * height_ratio
            gt_center_joints = torch.cat((center_x_joint, center_y_joint), dim=-1)
            for j,ct_joint in enumerate(gt_center_joints):
                ctx_int,cty_int = ct_joint.int()
                ctx,cty = ct_joint
                obj_target[batch_id,:,cty_int,ctx_int] = 1
                joint_offset_target[batch_id,0,cty_int,ctx_int] = ctx-ctx_int
                joint_offset_target[batch_id,1,cty_int,ctx_int] = cty-cty_int
        avg_factor1 = max(1,center_heatmap_target.eq(1).sum())
        avg_factor2 = max(1,obj_target.eq(1).sum())
        target_resullt = dict(
            center_heatmap_target = center_heatmap_target,
            slot_offset_target = slot_offset_target,
            slot_ori_target = slot_ori_target,
            slot_ct_target = slot_ct_target,
            obj_target = obj_target,
            joint_offset_target = joint_offset_target,
            avg_factor1 = avg_factor1,
            avg_factor2 = avg_factor2,
            slot_weight_target = slot_weight_target
        )
        return target_resullt

                
